"""
Plotting and visualization utilities for RINN implementation.
"""

import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
import numpy as np
import torch
from mpl_toolkits.mplot3d import Axes3D


def plot_2d_solution(points, values, title="Solution", figsize=(10, 8), 
                    scatter=False, levels=20, save_path=None):
    """
    Plot 2D solution using contour or scatter plot.
    
    Args:
        points: Points array (N, 2)
        values: Solution values (N,)
        title: Plot title
        figsize: Figure size
        scatter: Use scatter plot instead of contour
        levels: Number of contour levels
        save_path: Path to save figure
    """
    if isinstance(points, torch.Tensor):
        points = points.detach().cpu().numpy()
    if isinstance(values, torch.Tensor):
        values = values.detach().cpu().numpy().flatten()
    
    x = points[:, 0]
    y = points[:, 1]
    
    fig, ax = plt.subplots(figsize=figsize)
    
    if scatter:
        scatter = ax.scatter(x, y, c=values, cmap='viridis')
        plt.colorbar(scatter)
    else:
        # Create regular grid for contour plotting
        xi = np.linspace(x.min(), x.max(), 100)
        yi = np.linspace(y.min(), y.max(), 100)
        xi, yi = np.meshgrid(xi, yi)
        
        from scipy.interpolate import griddata
        zi = griddata((x, y), values, (xi, yi), method='linear')
        
        contour = ax.contourf(xi, yi, zi, levels=levels, cmap='viridis')
        plt.colorbar(contour)
    
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title(title)
    ax.set_aspect('equal')
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {save_path}")
    else:
        default_path = f"plot_{title.lower().replace(' ', '_')}.png"
        plt.savefig(default_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {default_path}")
    
    plt.close()


def plot_3d_solution(points, values, title="3D Solution", figsize=(12, 9), save_path=None):
    """
    Plot 3D surface of the solution.
    
    Args:
        points: Points array (N, 2)
        values: Solution values (N,)
        title: Plot title
        figsize: Figure size
        save_path: Path to save figure
    """
    if isinstance(points, torch.Tensor):
        points = points.detach().cpu().numpy()
    if isinstance(values, torch.Tensor):
        values = values.detach().cpu().numpy().flatten()
    
    x = points[:, 0]
    y = points[:, 1]
    
    # Create regular grid
    xi = np.linspace(x.min(), x.max(), 50)
    yi = np.linspace(y.min(), y.max(), 50)
    xi, yi = np.meshgrid(xi, yi)
    
    from scipy.interpolate import griddata
    zi = griddata((x, y), values, (xi, yi), method='linear')
    
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111, projection='3d')
    
    surf = ax.plot_surface(xi, yi, zi, cmap='viridis', alpha=0.7)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('u(x,y)')
    ax.set_title(title)
    
    fig.colorbar(surf)
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {save_path}")
    else:
        default_path = f"plot_{title.lower().replace(' ', '_')}.png"
        plt.savefig(default_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {default_path}")
    
    plt.close()


def plot_error_distribution(points, true_values, pred_values, title="Error Distribution", 
                          figsize=(12, 5), save_path=None):
    """
    Plot error distribution with both absolute error and histogram.
    
    Args:
        points: Points array (N, 2)
        true_values: True solution values (N,)
        pred_values: Predicted solution values (N,)
        title: Plot title
        figsize: Figure size
        save_path: Path to save figure
    """
    if isinstance(points, torch.Tensor):
        points = points.detach().cpu().numpy()
    if isinstance(true_values, torch.Tensor):
        true_values = true_values.detach().cpu().numpy().flatten()
    if isinstance(pred_values, torch.Tensor):
        pred_values = pred_values.detach().cpu().numpy().flatten()
    
    error = np.abs(true_values - pred_values)
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
    
    # Spatial error distribution
    scatter = ax1.scatter(points[:, 0], points[:, 1], c=error, cmap='Reds')
    ax1.set_xlabel('x')
    ax1.set_ylabel('y')
    ax1.set_title('Absolute Error Distribution')
    ax1.set_aspect('equal')
    plt.colorbar(scatter, ax=ax1)
    
    # Error histogram
    ax2.hist(error, bins=50, alpha=0.7, color='red', edgecolor='black')
    ax2.set_xlabel('Absolute Error')
    ax2.set_ylabel('Frequency')
    ax2.set_title('Error Histogram')
    ax2.grid(True, alpha=0.3)
    
    plt.suptitle(title)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {save_path}")
    else:
        default_path = f"plot_{title.lower().replace(' ', '_')}.png"
        plt.savefig(default_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {default_path}")
    
    plt.close()


def plot_training_history(history, metrics=None, figsize=(12, 8), save_path=None):
    """
    Plot training history for multiple metrics.
    
    Args:
        history: Dictionary with metric histories
        metrics: List of metrics to plot (if None, plot all)
        figsize: Figure size
        save_path: Path to save figure
    """
    if metrics is None:
        metrics = list(history.keys())
    
    n_metrics = len(metrics)
    if n_metrics <= 2:
        fig, axes = plt.subplots(1, n_metrics, figsize=figsize)
        if n_metrics == 1:
            axes = [axes]
    else:
        rows = (n_metrics + 1) // 2
        fig, axes = plt.subplots(rows, 2, figsize=figsize)
        axes = axes.flatten()
    
    for i, metric in enumerate(metrics):
        if metric in history:
            epochs = [e for e, _ in history[metric]]
            values = [v for _, v in history[metric]]
            
            axes[i].plot(epochs, values, 'b-', linewidth=2)
            axes[i].set_xlabel('Epoch')
            axes[i].set_ylabel(metric)
            axes[i].set_title(f'Training {metric}')
            axes[i].grid(True, alpha=0.3)
            axes[i].set_yscale('log' if 'loss' in metric.lower() or 'error' in metric.lower() else 'linear')
    
    # Hide unused subplots
    for i in range(len(metrics), len(axes)):
        axes[i].set_visible(False)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {save_path}")
    else:
        default_path = f"plot_{title.lower().replace(' ', '_')}.png"
        plt.savefig(default_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {default_path}")
    
    plt.close()


def plot_comparison(points, true_values, pred_values, title="True vs Predicted", 
                   figsize=(15, 5), save_path=None):
    """
    Plot comparison between true and predicted solutions.
    
    Args:
        points: Points array (N, 2)
        true_values: True solution values (N,)
        pred_values: Predicted solution values (N,)
        title: Plot title
        figsize: Figure size
        save_path: Path to save figure
    """
    fig, axes = plt.subplots(1, 3, figsize=figsize)
    
    # True solution
    plot_2d_solution_on_axis(axes[0], points, true_values, "True Solution")
    
    # Predicted solution
    plot_2d_solution_on_axis(axes[1], points, pred_values, "Predicted Solution")
    
    # Error
    if isinstance(true_values, torch.Tensor):
        true_values = true_values.detach().cpu().numpy().flatten()
    if isinstance(pred_values, torch.Tensor):
        pred_values = pred_values.detach().cpu().numpy().flatten()
    
    error = np.abs(true_values - pred_values)
    plot_2d_solution_on_axis(axes[2], points, error, "Absolute Error", cmap='Reds')
    
    plt.suptitle(title)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {save_path}")
    else:
        default_path = f"plot_{title.lower().replace(' ', '_')}.png"
        plt.savefig(default_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {default_path}")
    
    plt.close()


def plot_2d_solution_on_axis(ax, points, values, title, cmap='viridis', levels=20):
    """Helper function to plot 2D solution on existing axis."""
    if isinstance(points, torch.Tensor):
        points = points.detach().cpu().numpy()
    if isinstance(values, torch.Tensor):
        values = values.detach().cpu().numpy().flatten()
    
    x = points[:, 0]
    y = points[:, 1]
    
    # Create regular grid for contour plotting
    xi = np.linspace(x.min(), x.max(), 100)
    yi = np.linspace(y.min(), y.max(), 100)
    xi, yi = np.meshgrid(xi, yi)
    
    from scipy.interpolate import griddata
    zi = griddata((x, y), values, (xi, yi), method='linear')
    
    contour = ax.contourf(xi, yi, zi, levels=levels, cmap=cmap)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title(title)
    ax.set_aspect('equal')
    
    return contour


def plot_convergence_study(results_dict, title="Convergence Study", figsize=(10, 6), save_path=None):
    """
    Plot convergence study results.
    
    Args:
        results_dict: Dictionary with method names as keys and (n_points, errors) as values
        title: Plot title
        figsize: Figure size
        save_path: Path to save figure
    """
    fig, ax = plt.subplots(figsize=figsize)
    
    for method_name, (n_points, errors) in results_dict.items():
        ax.loglog(n_points, errors, 'o-', label=method_name, linewidth=2, markersize=6)
    
    ax.set_xlabel('Number of Points')
    ax.set_ylabel('Relative L2 Error')
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.legend()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {save_path}")
    else:
        default_path = f"plot_{title.lower().replace(' ', '_')}.png"
        plt.savefig(default_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {default_path}")
    
    plt.close()


def plot_basis_functions(network, domain, n_basis=None, figsize=(12, 8), save_path=None):
    """
    Visualize neural network basis functions.
    
    Args:
        network: Neural network with basis functions
        domain: Domain dictionary
        n_basis: Number of basis functions to plot (if None, plot all)
        figsize: Figure size
        save_path: Path to save figure
    """
    # Generate test points
    x = np.linspace(domain['x_min'], domain['x_max'], 50)
    y = np.linspace(domain['y_min'], domain['y_max'], 50)
    xx, yy = np.meshgrid(x, y)
    test_points = torch.tensor(np.column_stack([xx.flatten(), yy.flatten()]), dtype=torch.float32)
    
    # Get basis function outputs
    with torch.no_grad():
        if hasattr(network, 'get_basis_functions'):
            basis_outputs = network.get_basis_functions(test_points)
        else:
            # For standard networks, use hidden layer outputs
            basis_outputs = network[:-1](test_points)  # Exclude final layer
    
    if n_basis is None:
        n_basis = min(basis_outputs.shape[1], 16)
    
    cols = 4
    rows = (n_basis + cols - 1) // cols
    
    fig, axes = plt.subplots(rows, cols, figsize=figsize)
    if rows == 1:
        axes = axes.reshape(1, -1)
    
    for i in range(n_basis):
        row, col = i // cols, i % cols
        values = basis_outputs[:, i].numpy()
        
        # Reshape for contour plot
        values_2d = values.reshape(xx.shape)
        
        contour = axes[row, col].contourf(xx, yy, values_2d, levels=20, cmap='viridis')
        axes[row, col].set_title(f'Basis Function {i+1}')
        axes[row, col].set_aspect('equal')
    
    # Hide unused subplots
    for i in range(n_basis, rows * cols):
        row, col = i // cols, i % cols
        axes[row, col].set_visible(False)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {save_path}")
    else:
        default_path = f"plot_{title.lower().replace(' ', '_')}.png"
        plt.savefig(default_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {default_path}")
    
    plt.close()


def save_all_plots(results, experiment_name, save_dir="plots"):
    """Save all plots for an experiment."""
    import os
    os.makedirs(save_dir, exist_ok=True)
    
    base_path = os.path.join(save_dir, experiment_name)
    
    # Plot solution
    if 'points' in results and 'solution' in results:
        plot_2d_solution(results['points'], results['solution'], 
                        title=f"{experiment_name} - Solution",
                        save_path=f"{base_path}_solution.png")
    
    # Plot error distribution
    if all(k in results for k in ['points', 'true_solution', 'predicted_solution']):
        plot_error_distribution(results['points'], results['true_solution'], 
                              results['predicted_solution'],
                              title=f"{experiment_name} - Error Distribution",
                              save_path=f"{base_path}_error.png")
    
    # Plot training history
    if 'training_history' in results:
        plot_training_history(results['training_history'],
                            title=f"{experiment_name} - Training History",
                            save_path=f"{base_path}_training.png")